//! Mid-level translation step that translates type layouts to IR.

use thiserror::Error;

use crate::frontend::FieldType;
use crate::identifier::IdentString;
use crate::ir;
use crate::layout::layout::{Field, FieldOffset, ModuleLayout, StencilLayout, StructLayout};
use crate::mir::{Expr, TypedExpr};
use crate::ssa::{self, AssignmentTarget};
use crate::tysys::{
    typecheck_expr, ArgRef, BinOp, FieldRef, FnContext, GenericContext, IntEncoding, LiteralValue,
    LocalRef, ModuleContext, TypeError, ValueType,
};

#[derive(Debug, Error)]
pub enum Error {
    #[error("unsupported field type {0:?}")]
    UnsupportedFieldType(FieldType),

    #[error("unsupported array field type: {0:?}")]
    UnsupportedArrayFieldType(FieldType),

    #[error("typecheck: {0}")]
    Typecheck(#[from] TypeError),

    #[error("not yet implemented")]
    Unimplemented,

    #[error("{0}")]
    Other(String),
}

pub fn gen_ir_from_layout(ml: &ModuleLayout, name: &str) -> Result<ir::Module, Error> {
    let mod_str = IdentString::new_from_snake(name.to_owned()).expect("layout: mod name");
    let mut mb = ir::ModuleBuilder::new(mod_str);

    for (name, _id, stl) in ml.iter_tys() {
        let istr = IdentString::new_from_snake(name.clone()).expect("layout: struct name");
        match stl {
            StencilLayout::Struct(sl) => {
                let dt = gen_struct_data_type(sl, istr, ml)?;
                mb.add_data_type(dt);
            }
        }
    }

    Ok(mb.build())
}

pub fn gen_struct_data_type(
    sl: &StructLayout,
    name: IdentString,
    ml: &ModuleLayout,
) -> Result<ir::DataType, Error> {
    let mut dtb = ir::DataTypeBuilder::new(name);

    let ctor_ctx = StructFnContext {
        fields: vec![ValueType::ByteArray(sl.size())],
        args: vec![ValueType::ByteSlice],
        buf_field: FieldRef::from(0 as u16),
    };

    let ctor_buf_expr = Expr::new_arg(ArgRef::from(0));

    // Generate the size check thunk and set the size hint.
    let check_thunk = gen_struct_size_check(ctor_buf_expr, sl.size(), &ctor_ctx)?;
    dtb.set_check_thunk(check_thunk);
    dtb.set_size_hint(sl.size() as u64);

    // Generate each of the function accessors.
    for (name, field) in sl.fields() {
        let ctx = StructFnContext {
            fields: vec![ValueType::ByteArray(sl.size())],
            // FIXME fill this in
            args: Vec::new(),
            buf_field: FieldRef::from(0 as u16),
        };

        gen_struct_accessors(field, name, sl, ml, &mut dtb, &ctx)?;
    }

    Ok(dtb.build())
}

/// Messy context type for generating structs within a module.  This will be
/// simplified in the near future.
pub struct StructFnContext {
    /// The fields in the container.
    fields: Vec<ValueType>,

    /// Arguments.
    args: Vec<ValueType>,

    /// Buf field that we're wrapping.
    buf_field: FieldRef,
}

impl StructFnContext {
    /// Typechecks an expression, returning a type-tagged expression.
    pub fn check_expr(&self, expr: &Expr) -> Result<TypedExpr, TypeError> {
        typecheck_expr(&expr.to_int_tag(), self)
    }

    /// Gets the buffer field.
    pub fn buf_field(&self) -> FieldRef {
        self.buf_field
    }

    /// Returns an expr that accessor fns can use for the accesor
    pub fn get_acc_buf_expr(&self) -> Expr {
        Expr::new_field(self.buf_field)
    }
}

impl FnContext for StructFnContext {
    fn is_container(&self) -> bool {
        true
    }

    fn num_args(&self) -> u16 {
        self.args.len() as u16
    }

    fn get_arg_type(&self, arg: ArgRef) -> Option<ValueType> {
        self.args.get(arg.inner() as usize).cloned()
    }

    fn get_instance_field_type(&self, field: FieldRef) -> Option<ValueType> {
        self.fields.get(field.inner() as usize).cloned()
    }

    fn num_instance_fields(&self) -> u16 {
        self.fields.len() as u16
    }

    fn get_local_type(&self, _local: LocalRef) -> Option<ValueType> {
        // FIXME
        None
    }
}

impl ModuleContext for StructFnContext {}
impl GenericContext for StructFnContext {}

fn gen_struct_size_check(
    buf: Expr,
    size: usize,
    ctx: &StructFnContext,
) -> Result<ssa::BufThunk, Error> {
    // Construct expression to check if the size isn't what's the expected.
    let e1 = Expr::BufferLen((), Box::new(buf));
    let e2 = Expr::new_usize(size as u64);
    let e3 = Expr::new_ne(e1, e2);
    let et = ctx.check_expr(&e3).expect("layout: gen struct checks");

    let mut builder = ssa::SsaBuilder::new();
    let res = ssa::flatten_expr_rec(&et, &mut builder);
    builder.put_entry(ssa::SsaEntry::AbortIf(
        ssa::SsaArg::Local(res),
        ssa::AbortReason::BufMissized(size as u64),
    ));

    let thunk = builder.finalize_thunk();
    Ok(thunk)
}

fn gen_struct_accessors(
    f: &Field,
    fname: &str,
    sl: &StructLayout,
    ml: &ModuleLayout,
    builder: &mut ir::DataTypeBuilder,
    ctx: &StructFnContext,
) -> Result<(), Error> {
    // Generate the getter.
    let getter_name =
        IdentString::new_from_snake(format!("get_{fname}")).expect("modulegen: accessor name");
    let getter_fun = gen_getter_fun(f, getter_name, sl, ml, ctx)?;
    builder.add_function(getter_fun);

    // Generate the setter.
    let setter_name =
        IdentString::new_from_snake(format!("set_{fname}")).expect("modulegen: accessor name");
    let setter_fun = gen_setter_fun(f, setter_name, sl, ml, ctx)?;
    builder.add_function(setter_fun);

    Ok(())
}

fn gen_getter_fun(
    fdef: &Field,
    fun_name: IdentString,
    sl: &StructLayout,
    ml: &ModuleLayout,
    ctx: &StructFnContext,
) -> Result<ir::Function, Error> {
    let mut builder = ssa::SsaBuilder::new();
    let mut tc_arg_tys = Vec::new();
    let mut args = Vec::new();

    // For array types, we have to have index arguments corresponding to each
    // dimension it has.
    let nargs = fdef.ty().get_dimensions() as u16;
    for i in 0..nargs {
        let arg = ArgRef::from(i as u16);
        builder.set_arg_type(arg, ValueType::USize);
        tc_arg_tys.push(ValueType::USize);
        args.push(Expr::new_arg(arg));
    }

    // Compute the address of the thing we're looking for.
    let base_off_expr = gen_offset_expr(fdef.offset(), sl);
    let off_expr = match gen_array_index_offset_expr(fdef.ty(), &args, ml)? {
        Some(index_off) => Expr::new_binop(BinOp::Add, base_off_expr, index_off),
        None => base_off_expr,
    };

    // Construct the function for real.
    let fund_ty = get_fundamental_field_type(fdef.ty());

    gen_bounds_checks(fdef.ty(), &args, ml, &ctx, &mut builder)?;
    let buf_expr = ctx.get_acc_buf_expr();
    let getter_fun =
        gen_getter_fun_at(buf_expr, off_expr, fund_ty.clone(), &ctx, fun_name, builder)?;

    Ok(getter_fun)
}

fn gen_setter_fun(
    fdef: &Field,
    fun_name: IdentString,
    sl: &StructLayout,
    ml: &ModuleLayout,
    ctx: &StructFnContext,
) -> Result<ir::Function, Error> {
    let mut builder = ssa::SsaBuilder::new();
    let mut tc_arg_tys = Vec::new();
    let mut args = Vec::new();

    // For array types, we have to have index arguments corresponding to each
    // dimension it has.
    let nargs = fdef.ty().get_dimensions() as u16;
    for i in 0..nargs {
        let arg = ArgRef::from(i as u16);
        builder.set_arg_type(arg, ValueType::USize);
        tc_arg_tys.push(ValueType::USize);
        args.push(Expr::new_arg(arg));
    }

    // Compute the address of the thing we're looking for.
    let base_off_expr = gen_offset_expr(fdef.offset(), sl);
    let off_expr = match gen_array_index_offset_expr(fdef.ty(), &args, ml)? {
        Some(index_off) => Expr::new_binop(BinOp::Add, base_off_expr, index_off),
        None => base_off_expr,
    };

    // Construct the function for real.
    let fund_ty = get_fundamental_field_type(fdef.ty());
    let arg_ty = get_field_setter_arg_type(fund_ty);
    let arg = ArgRef::from(tc_arg_tys.len() as u16);
    builder.set_arg_type(arg, arg_ty.clone());
    tc_arg_tys.push(arg_ty);

    gen_bounds_checks(fdef.ty(), &args, ml, &ctx, &mut builder)?;
    let setter_fun = gen_setter_fun_at(off_expr, arg, fund_ty.clone(), &ctx, fun_name, builder)?;

    Ok(setter_fun)
}

/// Generates an expression that evaluates the byte offset of a series of
/// indexes into a nested array
/// structure.
fn gen_array_index_offset_expr(
    ty: &FieldType,
    arg_vals: &[Expr],
    ml: &ModuleLayout,
) -> Result<Option<Expr>, Error> {
    match ty {
        FieldType::Array(_size, aty) => {
            let layer_arg = arg_vals[0].clone();
            let stride_len = ml
                .compute_field_type_size(&aty)
                .ok_or_else(|| Error::UnsupportedFieldType(ty.clone()))?;

            // Compute this part of the stride.
            let stride_expr = Expr::new_usize(stride_len as u64);
            let comp_expr = Expr::new_binop(BinOp::Mul, layer_arg, stride_expr);

            // Now call recursively in case we have nested array structures.
            let expr = match gen_array_index_offset_expr(&aty, &arg_vals[1..], ml)? {
                Some(rest_expr) => Expr::new_binop(BinOp::Add, comp_expr, rest_expr),
                None => comp_expr,
            };

            Ok(Some(expr))
        }

        _ => Ok(None),
    }
}

/// Inserts bounds check entries in an SSA builder for interacting with an
/// array field.
fn gen_bounds_checks(
    ty: &FieldType,
    arg_vals: &[Expr],
    ml: &ModuleLayout,
    ctx: &StructFnContext,
    builder: &mut ssa::SsaBuilder,
) -> Result<(), Error> {
    match ty {
        FieldType::Array(size, aty) => {
            let arg_expr = arg_vals[0].clone();
            let dim_len_expr = Expr::new_usize(*size as u64);
            let check_expr = Expr::new_binop(BinOp::Ge, arg_expr, dim_len_expr);
            let check_expr_typed = ctx.check_expr(&check_expr)?;
            let check_vr = ssa::flatten_expr_rec(&check_expr_typed, builder);
            let abort_entry =
                ssa::SsaEntry::AbortIf(ssa::SsaArg::Local(check_vr), ssa::AbortReason::Oob);
            builder.put_entry(abort_entry);

            // Next layer.
            gen_bounds_checks(&aty, &arg_vals[1..], ml, ctx, builder)?;
            Ok(())
        }

        // No work to do.
        _ => Ok(()),
    }
}

/// Generates the rest of a function using an already partly set up.
fn gen_getter_fun_at(
    buf_expr: Expr,
    off_expr: Expr,
    fty: FieldType,
    ctx: &StructFnContext,
    fun_name: IdentString,
    mut builder: ssa::SsaBuilder,
) -> Result<ir::Function, Error> {
    let fun = match fty {
        FieldType::Bool => {
            let read_expr = Expr::new_read_int(buf_expr, off_expr, IntEncoding::byte());
            let cast_expr = Expr::new_cast(ValueType::Bool, read_expr);
            let read_expr_typed = ctx.check_expr(&cast_expr)?;
            let val_vr = ssa::flatten_expr_rec(&read_expr_typed, &mut builder);

            builder.finalize_function(Some(ssa::SsaArg::Local(val_vr)))
        }

        FieldType::Int(ity) => {
            let ienc = ity.to_big_encoding();
            let read_expr = Expr::new_read_int(buf_expr, off_expr, ienc);
            let read_expr_typed = ctx.check_expr(&read_expr)?;
            let val_vr = ssa::flatten_expr_rec(&read_expr_typed, &mut builder);

            builder.finalize_function(Some(ssa::SsaArg::Local(val_vr)))
        }

        FieldType::Buf(size) => {
            let size_expr = Expr::new_usize(size as u64);
            let read_expr = Expr::BufferRead(
                (),
                Box::new(buf_expr),
                Box::new(off_expr),
                Box::new(size_expr),
            );

            let read_expr_typed = ctx.check_expr(&read_expr)?;
            let val_vr = ssa::flatten_expr_rec(&read_expr_typed, &mut builder);

            builder.finalize_function(Some(ssa::SsaArg::Local(val_vr)))
        }

        FieldType::Struct(_) => {
            return Err(Error::Unimplemented);
        }

        _ => panic!("modulegen: tried to directly generate getter for type {fty:?}"),
    };

    Ok(ir::Function::new(fun_name, fun))
}

fn gen_setter_fun_at(
    off_expr: Expr,
    val_arg: ArgRef,
    fty: FieldType,
    ctx: &StructFnContext,
    fun_name: IdentString,
    mut builder: ssa::SsaBuilder,
) -> Result<ir::Function, Error> {
    let field_tgt = AssignmentTarget::Field(ctx.buf_field());

    let fun = match fty {
        FieldType::Bool => {
            let off_expr_typed = ctx.check_expr(&off_expr)?;
            let off_vr = ssa::flatten_expr_rec(&off_expr_typed, &mut builder);
            let val_entry = ssa::SsaExpr::Ternary(
                ssa::SsaArg::Arg(val_arg),
                ssa::SsaArg::Literal(LiteralValue::U8(1)),
                ssa::SsaArg::Literal(LiteralValue::U8(0)),
            );

            let val_vr = builder.put_expr_assign(val_entry, ValueType::U8);
            let set_entry = ssa::SsaEntry::WriteBuf(
                field_tgt,
                ssa::SsaArg::Local(off_vr),
                ssa::SsaArg::Local(val_vr),
            );

            builder.put_entry(set_entry);

            builder.finalize_function(None)
        }

        FieldType::Int(_) => {
            let off_expr_typed = ctx.check_expr(&off_expr)?;
            let off_vr = ssa::flatten_expr_rec(&off_expr_typed, &mut builder);
            let set_entry = ssa::SsaEntry::WriteBuf(
                field_tgt,
                ssa::SsaArg::Local(off_vr),
                ssa::SsaArg::Arg(val_arg),
            );

            builder.put_entry(set_entry);
            builder.finalize_function(None)
        }

        FieldType::Buf(_) => {
            let off_expr_typed = ctx.check_expr(&off_expr)?;
            let off_vr = ssa::flatten_expr_rec(&off_expr_typed, &mut builder);
            let set_entry = ssa::SsaEntry::WriteBuf(
                field_tgt,
                ssa::SsaArg::Local(off_vr),
                ssa::SsaArg::Arg(val_arg),
            );

            builder.put_entry(set_entry);
            builder.finalize_function(None)
        }

        FieldType::Struct(_) => {
            return Err(Error::Unimplemented);
        }

        _ => unimplemented!(),
    };

    Ok(ir::Function::new(fun_name, fun))
}

/// Generates an expression that computes the offset of a field, using whatever
/// other type data it needs.
fn gen_offset_expr(off: &FieldOffset, _container: &StructLayout) -> Expr {
    match off {
        FieldOffset::Fixed(off) => Expr::new_usize(*off),
    }
}

/// Gets the field type, stripping away any layers of arrays.
fn get_fundamental_field_type(ty: &FieldType) -> &FieldType {
    match ty {
        FieldType::Array(_, aty) => get_fundamental_field_type(aty),
        _ => ty,
    }
}

fn get_field_setter_arg_type(ty: &FieldType) -> ValueType {
    match ty {
        FieldType::Void => ValueType::Void,
        FieldType::Bool => ValueType::Bool,
        FieldType::Int(ity) => (*ity).into(),
        FieldType::Buf(size) => ValueType::ByteArray(*size),
        FieldType::Struct(id) => ValueType::ObjectInst(*id),
        _ => panic!("modulegen: type unable to be direct accessor type: {ty:?}"),
    }
}

#[cfg(test)]
mod tests {
    use crate::codegen::util::CodeBuffer;
    use crate::codegen::{python::PythonCg, traits::*};
    use crate::frontend::*;
    use crate::layout::*;

    use crate::tysys::{IntType, UserTypeId};

    use super::*;

    #[test]
    fn test_layout_cg_simple() {
        let f1 = Field::new(FieldOffset::Fixed(0), IntType::U64.into());
        let f2 = Field::new(FieldOffset::Fixed(8), IntType::U8.into());
        let layout = StructLayout::new(vec![("foo".to_owned(), f1), ("bar".to_owned(), f2)], 9);
        let ll = StencilLayout::Struct(layout.clone());
        let ml = ModuleLayout::new(vec![("world".to_owned(), UserTypeId::from_raw(42), ll)]);

        let mod_name = IdentString::new_from_snake("hello".to_owned()).unwrap();
        let mut mb = ir::ModuleBuilder::new(mod_name);

        let dt_name = IdentString::new_from_snake("world".to_owned()).unwrap();
        let dt = gen_struct_data_type(&layout, dt_name, &ml).expect("test: gen layout");

        mb.add_data_type(dt);
        let module = mb.build();

        let mut cgbuf = CodeBuffer::new(PythonCg.get_indent());
        let lir = PythonCg
            .convert_cir_to_lir(&module)
            .expect("test: specialize");
        PythonCg
            .generate_lang_ir(&lir, &mut cgbuf)
            .expect("test: codegen");

        let output = cgbuf.finalize();
        eprintln!("{output}");
    }

    #[test]
    fn test_layout_gen_simple() {
        let f1 = StructFieldDef::new("foo".into(), IntType::U32.into());
        let f2 = StructFieldDef::new("bar".into(), IntType::I64.into());
        let f3 = StructFieldDef::new("baz".into(), FieldType::Bool);
        let f4 = StructFieldDef::new("quux".into(), FieldType::Buf(7));
        let f5 = StructFieldDef::new("quack".into(), FieldType::int_array(5, IntType::I16));
        let fields = [f1, f2, f3, f4, f5]
            .into_iter()
            .map(StructEntry::Field)
            .collect::<Vec<_>>();
        let entry = Entry::Struct(UserTypeId::from_raw(1), StructDef::new(fields));
        let m = ModuleDef::new(vec![("Hello".to_owned(), entry)]);

        let ml = gen_mod_layout_from_mod_defn(&m).expect("test: gen mod layout");
        eprintln!("{:#?}", ml);

        let StencilLayout::Struct(sl) = ml.find_layout_by_name("Hello").unwrap();

        let mod_name = IdentString::from_parts(["hello"]).unwrap();
        let mut mb = ir::ModuleBuilder::new(mod_name);

        let dt_name = IdentString::from_parts(["world"]).unwrap();
        let dt = gen_struct_data_type(&sl, dt_name, &ml).expect("test: gen layout");

        mb.add_data_type(dt);
        let module = mb.build();

        let mut cgbuf = CodeBuffer::new(PythonCg.get_indent());
        let lir = PythonCg
            .convert_cir_to_lir(&module)
            .expect("test: specialize");
        PythonCg
            .generate_lang_ir(&lir, &mut cgbuf)
            .expect("test: codegen");

        let output = cgbuf.finalize();
        eprintln!("{output}");
    }
}
